import pandas as pd
from typing import Dict, List, Any, Optional, Union, Tuple
from sqlalchemy import text, inspect
from loguru import logger
import time

from .connection import get_db_connection

class DatabaseOperations:
    """数据库操作类，封装常见的数据库操作。"""
    
    def __init__(self):
        self.db = get_db_connection()
        self._ensure_connection()
        self.inspector = inspect(self.engine)
    
    def _ensure_connection(self):
        """确保数据库连接已建立"""
        try:
            # 尝试连接数据库
            self.engine = self.db.connect()
        except Exception as e:
            logger.error(f"数据库连接失败，尝试重新连接: {str(e)}")
            # 等待1秒后重试
            time.sleep(1)
            self.engine = self.db.connect()
    
    def get_tables(self) -> List[str]:
        """获取数据库中的所有表"""
        # 确保连接有效
        self._ensure_connection()
        # 刷新inspector以获取最新的数据库结构
        self.inspector = inspect(self.engine)
        return self.inspector.get_table_names()
    
    def get_table_schema(self, table_name: str) -> List[Dict[str, Any]]:
        """获取表结构信息"""
        columns = self.inspector.get_columns(table_name)
        return [{"name": col["name"], "type": str(col["type"]), 
                 "nullable": col["nullable"], "default": col.get("default")} 
                for col in columns]
    
    def get_table_relationships(self, table_name: str) -> List[Dict[str, Any]]:
        """获取表的外键关系"""
        fks = self.inspector.get_foreign_keys(table_name)
        return [{"constrained_columns": fk["constrained_columns"],
                 "referred_table": fk["referred_table"],
                 "referred_columns": fk["referred_columns"]}
                for fk in fks]
    
    def get_db_metadata(self) -> Dict[str, Any]:
        """获取数据库元数据，包括表名、结构等"""
        tables = self.get_tables()
        metadata = {}
        
        for table in tables:
            metadata[table] = {
                "schema": self.get_table_schema(table),
                "relationships": self.get_table_relationships(table)
            }
        
        return metadata
    
    def execute_query(self, query: str, params: Optional[Dict[str, Any]] = None) -> Tuple[List[Dict[str, Any]], Optional[str]]:
        """执行SQL查询并返回结果和可能的错误信息"""
        try:
            with self.db.get_session() as session:
                # 打印SQL查询和参数，便于调试
                logger.debug(f"执行SQL查询: {query}")
                if params:
                    logger.debug(f"SQL参数: {params}")
                
                # 执行查询
                if params:
                    result = session.execute(text(query), params)
                else:
                    result = session.execute(text(query))
                
                # 安全处理结果集转换
                try:
                    # 处理COUNT或聚合查询特殊情况
                    if "COUNT(" in query.upper() or "SUM(" in query.upper() or "AVG(" in query.upper():
                        # 对于聚合查询，直接获取第一行
                        first_row = result.fetchone()
                        if first_row is None:
                            return [], None
                        
                        # 如果是Row对象，使用_mapping属性获取字典
                        if hasattr(first_row, '_mapping'):
                            row_dict = dict(first_row._mapping)
                        else:
                            # 尝试将Row对象转换为字典
                            try:
                                row_dict = dict(first_row)
                            except Exception as e:
                                logger.warning(f"无法将Row转换为字典: {str(e)}")
                                # 如果是单值结果，创建一个包含该值的字典
                                if len(first_row) == 1:
                                    row_dict = {"count": first_row[0]}
                                else:
                                    # 创建索引键的字典
                                    row_dict = {f"column_{i}": val for i, val in enumerate(first_row)}
                        
                        # 打印转换后的结果
                        logger.debug(f"转换后的查询结果: {row_dict}")
                        return [row_dict], None
                    
                    # 标准查询处理
                    rows = []
                    for row in result:
                        # 如果是Row对象，使用_mapping属性获取字典
                        if hasattr(row, '_mapping'):
                            row_dict = dict(row._mapping)
                        else:
                            # 尝试将Row对象转换为字典
                            try:
                                row_dict = dict(row)
                            except Exception as e:
                                logger.warning(f"无法将Row转换为字典: {str(e)}")
                                # 创建索引键的字典
                                row_dict = {f"column_{i}": val for i, val in enumerate(row)}
                        
                        rows.append(row_dict)
                    
                    # 打印查询结果的条数
                    logger.debug(f"查询返回了 {len(rows)} 条结果")
                    return rows, None
                    
                except Exception as e:
                    logger.error(f"处理查询结果时出错: {str(e)}")
                    logger.error(f"结果类型: {type(result)}")
                    # 尝试获取部分原始结果信息
                    try:
                        first_row = result.fetchone()
                        logger.error(f"第一行结果类型: {type(first_row)}")
                        logger.error(f"第一行结果内容: {first_row}")
                    except:
                        pass
                    
                    return [], f"处理查询结果失败: {str(e)}"
                
        except Exception as e:
            error_msg = f"查询执行失败: {str(e)}"
            logger.error(error_msg)
            return [], error_msg
    
    def execute_write_query(self, query: str, params: Optional[Dict[str, Any]] = None) -> Tuple[bool, Optional[str]]:
        """执行写入操作（INSERT/UPDATE/DELETE）并返回状态和可能的错误信息"""
        try:
            with self.db.get_session() as session:
                if params:
                    session.execute(text(query), params)
                else:
                    session.execute(text(query))
                session.commit()
                return True, None
        except Exception as e:
            error_msg = f"写入操作失败: {str(e)}"
            logger.error(error_msg)
            return False, error_msg
    
    def query_to_dataframe(self, query: str, params: Optional[Dict[str, Any]] = None) -> Tuple[Optional[pd.DataFrame], Optional[str]]:
        """执行查询并将结果转换为pandas DataFrame"""
        try:
            data, error = self.execute_query(query, params)
            if error:
                return None, error
            df = pd.DataFrame(data)
            return df, None
        except Exception as e:
            error_msg = f"查询转换为DataFrame失败: {str(e)}"
            logger.error(error_msg)
            return None, error_msg
    
    def get_sample_data(self, table_name: str, limit: int = 5) -> Tuple[Optional[pd.DataFrame], Optional[str]]:
        """获取表的样本数据"""
        query = f"SELECT * FROM {table_name} LIMIT {limit}"
        return self.query_to_dataframe(query)
    
    def count_rows(self, table_name: str, condition: str = "") -> int:
        """计算表中的行数，可以附加条件"""
        query = f"SELECT COUNT(*) as count FROM {table_name}"
        if condition:
            query += f" WHERE {condition}"
        result, error = self.execute_query(query)
        if error or not result:
            return 0
        return result[0]["count"]
    
    def close(self):
        """关闭数据库连接"""
        self.db.close()


def get_db_operations() -> DatabaseOperations:
    """获取数据库操作实例"""
    return DatabaseOperations() 